package main

import (
	"cs466/findmotif/fasta"
	"cs466/findmotif/gibbs"

	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"math"
	"os"
	"time"
	"runtime"
	"runtime/pprof"
	"encoding/csv"
)

var jobs = flag.Int("jobs", 2, "Number of Gibbs-Sampling jobs to run simultaneously")
var datadirname = flag.String("data", "data", "Path to directory containing benchmarks")
	
// profiling flags
var cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")

// http://en.wikipedia.org/wiki/Jensen%E2%80%93Shannon_divergence
// an always-defined relative-entropy-based distance between probability
// distributions.  (relative entropy is frequently undefined)
func JensenShannonDivergence(P [][]float64, Q [][]float64, ml float64) (ret float64) {
	if len(P) != len(Q) || len(P[0]) != len(Q[0]) {
		log.Panicf("PWM Size mismatch %d x %d, %d x %d", len(P), len(P[0]), len(Q), len(Q[0]))
	}

	M := make([][]float64, len(P))
	for i := range M {
		M[i] = make([]float64, len(P[i]))
		for j := range M[i] {
			M[i][j] = P[i][j]/2 + Q[i][j]/2
		}
	}

	return RelativeEntropyOf(P, M, ml)/2 + RelativeEntropyOf(Q, M, ml)/2
}

func RelativeEntropyOf(pwm1 [][]float64, pwm2 [][]float64, ml float64) (ret float64) {
	if len(pwm1) != len(pwm2) || len(pwm1[0]) != len(pwm2[0]) {
		log.Panicf("PWM Size mismatch %d x %d, %d x %d", len(pwm1), len(pwm1[0]), len(pwm2), len(pwm2[0]))
	}

	for i, col := range pwm1 {
		for j, el := range col {
			p := el / ml
			q := pwm2[i][j] / ml

			// relative entropy is undefined if not (q = 0 => p = 0)
			if p > 0 && q == 0 {
				return math.NaN()
			}

			// see the Wikipedia article: if p is 0, the limit of what we would
			// add here as p approaches 0 is 0, so don't add anything; if we
			// tried, we would evaluate ln(0), which is undefined.
			if p > 0 {
				ret += p * math.Log(p/q)
			}
		}
	}

	return ret
}

func EvaluateResults(dirname string) (relentropy float64, divergence float64, overlap int) {
	relentropy = 0
	overlap = 0
	// compute relative entropy
	pwm1_name := fmt.Sprintf("%s/motif.txt", dirname)
	pwm2_name := fmt.Sprintf("%s/predictedmotif.txt", dirname)
	ml_name := fmt.Sprintf("%s/motiflength.txt", dirname)

	ml, err := fasta.LoadMotifLength(ml_name)
	if err != nil {
		log.Fatal(err)
	}
	pwm1, err := fasta.LoadMotif(pwm1_name)
	if err != nil {
		log.Fatal(err)
	}
	pwm2, err := fasta.LoadMotif(pwm2_name)
	if err != nil {
		log.Fatal(err)
	}

	relentropy = RelativeEntropyOf(pwm1, pwm2, float64(ml))
	divergence = JensenShannonDivergence(pwm1, pwm2, float64(ml))

	// compute site overlap
	sites_name := fmt.Sprintf("%s/sites.txt", dirname)
	predsites_name := fmt.Sprintf("%s/predictedsites.txt", dirname)

	sites, err := fasta.LoadSites(sites_name)
	if err != nil {
		log.Fatal(err)
	}

	predsites, err := fasta.LoadSites(predsites_name)
	if err != nil {
		log.Fatal(err)
	}

	if len(sites) != len(predsites) {
		log.Panicf("Sites length mismatch %d, %d\n", len(sites), len(predsites))
	}

	for i := range sites {
		o := int(ml) - int(math.Abs(float64(sites[i]-predsites[i])))
		if o > 0 {
			overlap = overlap + o
		}
	}

	return
}

func main() {
	flag.Parse()
	runtime.GOMAXPROCS(*jobs)
	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	datasetsinfo, err := ioutil.ReadDir(*datadirname)
	if err != nil {
		log.Fatal(err)
	}

	datasets := make([]string, 0, len(datasetsinfo))

	for _, v := range datasetsinfo {
		if v.IsDir() {
			datasets = append(datasets, v.Name())
		}
	}

	runtimes := make([]float64, len(datasets))

	threads := make([]chan bool, *jobs)
	N := len(datasets)
	for i := 0; i < *jobs; i++ {
		threads[i] = make(chan bool)
		start, end := (i*N) / *jobs, ((i+1)*N) / *jobs

		go func(dirnames []string, times []float64, done chan bool, id int) {
			for i, v := range dirnames {
				datapath := fmt.Sprintf("%s/%s", *datadirname, v)
				start := time.Now()

				gibbs.Main(datapath)

				times[i] = time.Since(start).Seconds()

				log.Printf("(%d) %s:\t%f\n", id, datapath, times[i])
			}
			done <- true
		}(datasets[start:end], runtimes[start:end], threads[i], i)
	}

	for _, thread := range threads {
		<-thread
	}

	records := [][]string{}
	for i, v := range datasets {
		evalname := fmt.Sprintf("%s/%s.eval", *datadirname, v)
		f, err := os.Create(evalname)
		if err != nil {
			log.Fatal(err)
		}
		defer f.Close()

		dirname := fmt.Sprintf("%s/%s", *datadirname, v)
		relentropy, divergence, overlap := EvaluateResults(dirname)

		_, err = fmt.Fprintf(f, "Relative Entropy =\t%f\n", relentropy)
		if err != nil {
			log.Fatal(err)
		}
		_, err = fmt.Fprintf(f, "J-S Divergence =\t%f\n", divergence)
		if err != nil {
			log.Fatal(err)
		}
		_, err = fmt.Fprintf(f, "Site Overlap =\t%d\n", overlap)
		if err != nil {
			log.Fatal(err)
		}
		_, err = fmt.Fprintf(f, "Runtime =\t%f\n", runtimes[i])
		if err != nil {
			log.Fatal(err)
		}

		var icpc float64
		var ml, sc int
		_, err = fmt.Sscanf(v, "%f-%d-500-%d", &icpc, &ml, &sc)
		if err != nil {
			log.Fatal(err)
		}

		record := []string{
			fmt.Sprint(icpc),
			fmt.Sprint(ml),
			fmt.Sprint(sc),
			fmt.Sprint(divergence),
			fmt.Sprint(overlap),
			fmt.Sprint(runtimes[i]),
		}
		records = append(records, record)
	}

	csvname := fmt.Sprintf("%s/eval.csv", *datadirname)
	f, err := os.Create(csvname)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()

	err = csv.NewWriter(f).WriteAll(records)
	if err != nil {
		log.Fatal(err)
	}
}
